import java.util.ArrayList;
import java.util.List;

class BSTChallenges {

    // 找到给定节点的中序后继节点
    public TreeNode inorderSuccessor(TreeNode root, TreeNode p) {
        TreeNode successor = null; // 初始化后继节点为null
        while (root != null) { // 遍历BST
            if (p.val < root.val) { // 如果p的值小于当前节点的值
                successor = root; // 当前节点可能是后继节点
                root = root.left; // 继续在左子树中查找
            } else { // 如果p的值大于或等于当前节点的值
                root = root.right; // 继续在右子树中查找
            }
        }
        return successor; // 返回找到的后继节点
    }

    // 计算给定范围内所有节点值的和
    public int rangeSumBST(TreeNode root, int low, int high) {
        if (root == null) { // 如果当前节点为空，返回0
            return 0;
        }
        if (root.val < low) { // 如果当前节点的值小于low，只在右子树中查找
            return rangeSumBST(root.right, low, high);
        }
        if (root.val > high) { // 如果当前节点的值大于high，只在左子树中查找
            return rangeSumBST(root.left, low, high);
        }
        // 当前节点的值在范围内，加上当前节点的值，并递归计算左右子树
        return root.val + rangeSumBST(root.left, low, high) + rangeSumBST(root.right, low, high);
    }

    // 合并两个BST并返回排序后的所有元素列表
    public List<Integer> getAllElements(TreeNode root1, TreeNode root2) {
        List<Integer> list1 = new ArrayList<>(); // 存储第一个BST的中序遍历结果
        List<Integer> list2 = new ArrayList<>(); // 存储第二个BST的中序遍历结果
        inorderTraversal(root1, list1); // 对第一个BST进行中序遍历
        inorderTraversal(root2, list2); // 对第二个BST进行中序遍历
        return mergeSortedLists(list1, list2); // 合并两个有序列表并返回
    }

    // 中序遍历BST，将节点值存入列表
    private void inorderTraversal(TreeNode root, List<Integer> list) {
        if (root == null) { // 如果当前节点为空，返回
            return;
        }
        inorderTraversal(root.left, list); // 递归遍历左子树
        list.add(root.val); // 将当前节点的值加入列表
        inorderTraversal(root.right, list); // 递归遍历右子树
    }

    // 合并两个有序列表
    private List<Integer> mergeSortedLists(List<Integer> list1, List<Integer> list2) {
        List<Integer> mergedList = new ArrayList<>(); // 初始化合并后的列表
        int i = 0, j = 0; // 初始化两个指针
        while (i < list1.size() && j < list2.size()) { // 遍历两个列表
            if (list1.get(i) < list2.get(j)) { // 如果list1的当前元素小于list2的当前元素
                mergedList.add(list1.get(i++)); // 将list1的当前元素加入合并列表
            } else { // 否则
                mergedList.add(list2.get(j++)); // 将list2的当前元素加入合并列表
            }
        }
        while (i < list1.size()) { // 如果list1还有剩余元素
            mergedList.add(list1.get(i++)); // 将剩余元素加入合并列表
        }
        while (j < list2.size()) { // 如果list2还有剩余元素
            mergedList.add(list2.get(j++)); // 将剩余元素加入合并列表
        }
        return mergedList; // 返回合并后的列表
    }

    public static void main(String[] args) {
        BSTChallenges bst = new BSTChallenges(); // 创建BSTChallenges对象

        // 示例用法
        TreeNode root = new TreeNode(5); // 创建第一个BST的根节点
        root.left = new TreeNode(3); // 左子节点
        root.right = new TreeNode(7); // 右子节点
        root.left.left = new TreeNode(2); // 左子节点的左子节点
        root.left.right = new TreeNode(4); // 左子节点的右子节点
        root.right.left = new TreeNode(6); // 右子节点的左子节点
        root.right.right = new TreeNode(8); // 右子节点的右子节点

        TreeNode p = root.left.right; // 选择节点值为4的节点
        TreeNode successor = bst.inorderSuccessor(root, p); // 找到中序后继节点
        System.out.println("Inorder successor of " + p.val + " is " + (successor != null ? successor.val : "null")); // 输出后继节点

        int low = 3, high = 7; // 定义范围
        int sum = bst.rangeSumBST(root, low, high); // 计算范围内节点值的和
        System.out.println("Sum of nodes in range [" + low + ", " + high + "] is " + sum); // 输出结果

        TreeNode root2 = new TreeNode(1); // 创建第二个BST的根节点
        root2.right = new TreeNode(9); // 右子节点
        List<Integer> allElements = bst.getAllElements(root, root2); // 合并两个BST并返回排序后的列表
        System.out.println("All elements from both BSTs: " + allElements); // 输出合并后的列表
    }
}
